//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this sample source code is subject to the terms of the Microsoft
// license agreement under which you licensed this sample source code. If
// you did not accept the terms of the license agreement, you are not
// authorized to use this sample source code. For the terms of the license,
// please see the license agreement between you and Microsoft or, if applicable,
// see the LICENSE.RTF on your install media or the root of your tools installation.
// THE SAMPLE SOURCE CODE IS PROVIDED "AS IS", WITH NO WARRANTIES OR INDEMNITIES.
//
#include <windows.h>
#include <bootStoreUtils.h>
#include <fls.h>
#include <bootPart.h>

//------------------------------------------------------------------------------

#define PARTITIONS_IN_MBR           4

//------------------------------------------------------------------------------

BOOL
BootStoreIsThisMBR(
    UINT32 sectorSize,
    UINT8 *pMbrSector
    )
{
    BOOL rc = FALSE;

    // Check for MBR
    if ((*(UINT16*)&pMbrSector[BOOT_SIZE - 2] != BOOTSECTRAILSIGH) ||
        ((pMbrSector[0] != BS2BYTJMP) && (pMbrSector[0] != BS3BYTJMP)))
        goto cleanUp;

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreIsThisLayoutSector(
    UINT32 sectorSize,
    UINT8 *pLayoutSector
    )
{
    BOOL rc = FALSE;

    // Check for layout sector
    if (memcmp(pLayoutSector, FLASH_LAYOUT_SIG, FLASH_LAYOUT_SIG_SIZE) != 0)
        goto cleanUp;

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreGetReservedRegionInfo(
    UINT32 sectorSize,
    UINT8 *pLayoutSector,
    LPCSTR name,
    UINT32 *pStartSector,
    UINT32 *pSectors
    )
{
    BOOL rc = FALSE;
    FlashLayoutSector* pLayout = (FlashLayoutSector*)pLayoutSector;
    UINT8* pInfo = (UINT8*)&pLayout[1];
    FlashRegion* pRegion;
    ReservedEntry* pReservedRegion;
    UINT32 sectorsPerBlock;
    UINT32 reservedRegions, ix;


    // Check from layout sector
    if (memcmp(pLayoutSector, FLASH_LAYOUT_SIG, FLASH_LAYOUT_SIG_SIZE) != 0)
        goto cleanUp;

    // Get sectors per block, if there is no region assume 1 sector
    pRegion = (FlashRegion*)&pInfo[pLayout->cbReservedEntries];    
    if (pLayout->cbRegionEntries < sizeof(FlashRegion))
        {
        sectorsPerBlock = 1;
        }
    else
        {
        sectorsPerBlock = pRegion->dwSectorsPerBlock;
        }
    
    // Get number of reserved regions
    reservedRegions = pLayout->cbReservedEntries/sizeof(ReservedEntry);
    pReservedRegion = (ReservedEntry*)&pInfo[0];

    // Find reserved region by name
    for (ix = 0; ix < reservedRegions; ix++)
        {
        if (strcmp(pReservedRegion[ix].szName, name) == 0) break;
        }
    if (ix >= reservedRegions) goto cleanUp;

    // Copy information    
    if (pStartSector != NULL)
        {
        *pStartSector = pReservedRegion[ix].dwStartBlock * sectorsPerBlock;
        }
    if (pSectors != NULL)
        {
        *pSectors = pReservedRegion[ix].dwNumBlocks * sectorsPerBlock;
        }

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreGetReservedRegionName(
    UINT32 sectorSize,
    UINT8 *pLayoutSector,
    UINT32 index,
    LPSTR  pNameBuffer,
    UINT32 nameBufferSize
    )
{
    BOOL rc = FALSE;
    FlashLayoutSector* pLayout = (FlashLayoutSector*)pLayoutSector;
    UINT8* pInfo = (UINT8*)&pLayout[1];
    ReservedEntry* pReservedRegion;
    UINT32 reservedRegions;


    // Check from layout sector
    if (memcmp(pLayoutSector, FLASH_LAYOUT_SIG, FLASH_LAYOUT_SIG_SIZE) != 0)
        goto cleanUp;

    // Get number of reserved regions
    reservedRegions = pLayout->cbReservedEntries/sizeof(ReservedEntry);
    pReservedRegion = (ReservedEntry*)&pInfo[0];
    if (index >= reservedRegions) goto cleanUp;

    // Copy name
    PREFAST_SUPPRESS(419, "nameBufferSize is pNameBuffer size")
    strncpy(pNameBuffer, pReservedRegion[index].szName, nameBufferSize);
    pNameBuffer[nameBufferSize - 1] = '\0';

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreGetFlashRegionInfo(
    UINT32 sectorSize,
    UINT8 *pLayoutSector,
    UINT32 index,
    UINT32 *pSize,
    UINT32 *pCompactSize
    )
{
    BOOL rc = FALSE;
    FlashLayoutSector* pLayout = (FlashLayoutSector*)pLayoutSector;
    UINT8* pInfo = (UINT8*)&pLayout[1];
    FlashRegion* pRegion;
    UINT32 regions;


    // Check from layout sector
    if (memcmp(pLayoutSector, FLASH_LAYOUT_SIG, FLASH_LAYOUT_SIG_SIZE) != 0)
        goto cleanUp;

    // Get number of regions defined in layout sector
    regions = pLayout->cbRegionEntries/sizeof(FlashRegion);
    if (index >= regions) goto cleanUp;

    // Get region structure
    pRegion  = (FlashRegion*)&pInfo[pLayout->cbReservedEntries];
    pRegion += index;

    // Copy information    
    if (pSize != NULL)
        {
        if (pRegion->dwNumLogicalBlocks != -1)
            {
            *pSize = pRegion->dwBytesPerBlock * pRegion->dwNumLogicalBlocks;
            }
        else
            {
            *pSize = -1;
            }
        }
    if (pCompactSize != NULL)
        {
        *pCompactSize = pRegion->dwBytesPerBlock * pRegion->dwCompactBlocks;
        }

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreGetPartitionType(
    UINT32 sectorSize,
    UINT8 *pMbrSector,
    UINT32 index,
    UINT32 *pType
    )
{
    BOOL rc = FALSE;
    PARTENTRY partitionTable[PARTITIONS_IN_MBR];


    // Check for MBR
    if ((*(UINT16*)&pMbrSector[BOOT_SIZE - 2] != BOOTSECTRAILSIGH) ||
        ((pMbrSector[0] != BS2BYTJMP) && (pMbrSector[0] != BS3BYTJMP)))
        goto cleanUp;

    // We support only limited number of partitions
    if (index >= PARTITIONS_IN_MBR) goto cleanUp;
    
    // Copy partition table to avoid alignment issue
    memcpy(
        partitionTable, &pMbrSector[BOOT_SIZE - 2 - sizeof(partitionTable)],
        sizeof(partitionTable)
        );

    // Return type (zero is empty)
    *pType = partitionTable[index].Part_FileSystem;
    if (*pType == 0) goto cleanUp;

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreGetPartitionInfo(
    UINT32 sectorSize,
    UINT8 *pMbrSector,
    UINT32 type,
    UINT32 index,
    UINT32 *pStartSector,
    UINT32 *pSectors
    )
{
    BOOL rc = FALSE;
    PARTENTRY partitionTable[PARTITIONS_IN_MBR];
    UINT ix;


    // Check for MBR
    if ((*(UINT16*)&pMbrSector[BOOT_SIZE - 2] != BOOTSECTRAILSIGH) ||
        ((pMbrSector[0] != BS2BYTJMP) && (pMbrSector[0] != BS3BYTJMP)))
        goto cleanUp;

    // We support only limited number of partitions
    if (index >= PARTITIONS_IN_MBR) goto cleanUp;
    
    // Copy partition table to avoid alignment issue
    memcpy(
        partitionTable, &pMbrSector[BOOT_SIZE - 2 - sizeof(partitionTable)],
        sizeof(partitionTable)
        );

    ix = 0;
    while (ix < PARTITIONS_IN_MBR)
        {
        // Type zero denote empty slot
        if (partitionTable[ix].Part_FileSystem == type)
            {
            if (index == 0) break;
            index--;
            }
        ix++;
        }
    if (ix >= PARTITIONS_IN_MBR) goto cleanUp;

    // Copy info
    if (pStartSector != NULL)
        {
        *pStartSector = partitionTable[ix].Part_StartSector;
        }
    if (pSectors != NULL)
        {
        *pSectors = partitionTable[ix].Part_TotalSectors;
        }

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

BOOL
BootStoreUpdatePartitionInfo(
    UINT32 sectorSize,
    UINT8 *pMbrSector,
    UINT32 type,
    UINT32 index,
    UINT32 startSector,
    UINT32 sectors
    )
{
    BOOL rc = FALSE;
    PARTENTRY partitionTable[PARTITIONS_IN_MBR];
    UINT ix;


    // Check for MBR
    if ((*(UINT16*)&pMbrSector[BOOT_SIZE - 2] != BOOTSECTRAILSIGH) ||
        ((pMbrSector[0] != BS2BYTJMP) && (pMbrSector[0] != BS3BYTJMP)))
        goto cleanUp;

    // We support only limited number of partitions
    if (index >= PARTITIONS_IN_MBR) goto cleanUp;
    
    // Copy partition table to avoid alignment issue
    memcpy(
        partitionTable, &pMbrSector[BOOT_SIZE - 2 - sizeof(partitionTable)],
        sizeof(partitionTable)
        );

    ix = 0;
    while (ix < PARTITIONS_IN_MBR)
        {
        // Type zero denote empty slot
        if (partitionTable[ix].Part_FileSystem == type)
            {
            if (index == 0) break;
            index--;
            }
        ix++;
        }
    if (ix >= PARTITIONS_IN_MBR) goto cleanUp;

    // Update info
    partitionTable[ix].Part_StartSector = startSector;
    partitionTable[ix].Part_TotalSectors = sectors;

    // Done
    rc = TRUE;
    
cleanUp:
    return rc;
}

//------------------------------------------------------------------------------

